;; Non-robust row echelon form functions
;; Copyright (C) 2025  Rocks Mazama
;;
;; This program is free software: you can redistribute it and/or modify
;; it under the terms of the GNU General Public License as published by
;; the Free Software Foundation, either version 3 of the License, or
;; (at your option) any later version.
;;
;; This program is distributed in the hope that it will be useful,
;; but WITHOUT ANY WARRANTY; without even the implied warranty of
;; MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
;; GNU General Public License for more details.
;;
;; You should have received a copy of the GNU General Public License
;; along with this program.  If not, see <https://www.gnu.org/licenses/>.

(ns mazama.linalg.r-e-f)

;; See mazama.linalg.matrix for a description of the map used to describe
;; matrices.
;;
;; Augmented matrices can also be used. Augmented matrices of the form [A|b]
;; have a key :acols the value of which is the number of columns in the matrix
;; A. :ncols represents the number of columns in the entire augmented matrix.

(defn- get-better-pivot
  "When `row` contains a better pivot in column `j` than `pivot-value`."
  [j row pivot-value]
  (let [v (nth row j)]
    (when v
      (if pivot-value
        (when (> (abs v) (abs pivot-value))
          v)
        (when-not (zero? v)
          v)))))

(defn- ero-swap-rows-fn
  "Return a function to be used with map-indexed to swap values in positions `j` and `k`."
  [j vj k vk]
  (fn [i x]
    (cond
      (== i j) vk
      (== i k) vj
      :else x)))

(defn- find-pivot-info
  "When column `j` contains a pivot then return the pivot info."
  [j rows]
  (loop [rows (seq rows)
         i 0
         pivot-value nil
         pivot-row nil
         pivot-row-index nil]
    (if rows
      (let [v (get-better-pivot j (first rows) pivot-value)]
        (if v
          (recur (next rows) (inc i) v (first rows) i)
          (recur (next rows) (inc i) pivot-value pivot-row pivot-row-index)))
      (when pivot-value
        {:pivot-column-index j
         :pivot-value pivot-value
         :pivot-row pivot-row
         :pivot-row-index pivot-row-index}))))

(defn- pivotize-column
  "When column `j` contains a pivot then move the pivot row to the top."
  [j rows ero-scale]
  (when-let [info (find-pivot-info j rows)]
    (if (zero? (:pivot-row-index info))
      (assoc info :other-rows (next rows) :ero-scale ero-scale)
      (let [f (ero-swap-rows-fn 0
                                (first rows)
                                (:pivot-row-index info)
                                (:pivot-row info))
            rs (map-indexed f rows)]
        (assoc info :pivot-row (first rs) :other-rows (next rs) :ero-scale (- ero-scale))))))

(defn- pivotize
  "Beginning at column `j` configure the rows for the next pivot."
  [j rows jcols ero-scale]
  (loop [j j]
    (when (< j jcols)
      (let [info (pivotize-column j rows ero-scale)]
        (or info (recur (inc j)))))))

(defn- eliminate
  [{:keys [pivot-column-index pivot-value pivot-row other-rows] :as info}]
  (letfn [(f [row]
            (let [v (nth row pivot-column-index)
                  sf (/ v pivot-value)]
              (map #(- %1 (* sf %2)) row pivot-row)))]
    (assoc info :other-rows (map f other-rows))))

(defn- r-e-f-step
  [j rows jcols ero-scale]
  (when-let [info (pivotize j rows jcols ero-scale)]
    (eliminate info)))

(defn- data-as-rows
  "Return the matrix `data` as a sequence of rows, each row being a sequence of elements."
  [ncols data]
  (partition ncols data))

(defn r-e-f-helper
  "Compute a row echelon form of the matrix represented by `rows`."
  [rows jcols accumulator]  ; jcols is the number of columns to consider
    (loop [rows (seq rows)
           j 0                          ; Column index to start searching for pivots
           pivot-column-indexes accumulator
           pivot-rows accumulator
           zero-rows nil
           ero-scale 1]
      (if rows
        (if-let [{:keys [pivot-column-index pivot-row other-rows ero-scale]} (r-e-f-step j rows jcols ero-scale)]
          (recur (seq other-rows)
                 (inc pivot-column-index)
                 (conj pivot-column-indexes pivot-column-index)
                 (conj pivot-rows pivot-row)
                 zero-rows
                 ero-scale)
          (recur nil
                 j
                 pivot-column-indexes
                 pivot-rows
                 rows
                 ero-scale))
        {:pivot-column-indexes pivot-column-indexes
         :pivot-rows pivot-rows
         :zero-rows zero-rows
         :ero-scale ero-scale})))

(defn r-e-f
  "Compute a row echelon form of the matrix `m`.
  `m` can be an augmented matrix (with key :acols)."
  ([m]
   (r-e-f m []))
  ([{:keys [acols ncols data]} accumulator]
   (r-e-f (data-as-rows ncols data) (or acols ncols) accumulator))
  ([rows jcols accumulator]
   (r-e-f-helper rows jcols accumulator)))
